"""
dataset: MNIST

alg: a t-Distributed Stochastic Neighbour Embedding (t-SNE) decomposition
"""

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

from sklearn import datasets
from sklearn import manifold
if __name__ == '__main__':
    # Load data
    data = datasets.fetch_openml(
        'mnist_784',
        version=1,
        return_X_y=True
    )
    pixel_values, targets = data
    targets = targets.astype(int)

    # Check if data is loaded correctly
    if len(pixel_values) == 0:
        raise ValueError("Failed to load pixel values from MNIST dataset.")

    # Visualize a single image
    # Assuming we want to display the second image in the dataset
    single_image = pixel_values.values[1, :].reshape(28, 28)
    plt.imshow(single_image, cmap='gray')
    plt.title(f"Label: {targets[1]}")
    plt.axis('off')  # Hide axes for better visualization
    plt.show()

    tsne = manifold.TSNE(n_components=2, random_state=42)
    transformed_data = tsne.fit_transform(pixel_values.values[:3000, :])

    tsne_df = pd.DataFrame(
        np.column_stack((transformed_data, targets[:3000])),
        columns=["x", "y", "targets"]
    )
    tsne_df.loc[:, "targets"] = tsne_df.targets.astype(int)

    grid = sns.FacetGrid(tsne_df, hue="targets")
    grid.map(plt.scatter, "x", "y").add_legend()